#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
教室查询Agent示例
演示如何使用Langchain Agent调用教室查询工具
"""

import os

from dotenv import load_dotenv

load_dotenv()

from langchain_openai import ChatOpenAI
from langchain.agents import create_agent

from .login_system import BJTSULogin
from .classroom_tool_refactored import ClassroomQueryTool


def main():
    print("🎓 北京交通大学教室查询Agent演示")
    print("=" * 40)

    # 获取智谱AI API密钥（如果有的话）
    zhipu_api_key = os.getenv("ZHIPU_API_KEY")
    if not zhipu_api_key:
        print("⚠️  未配置智谱AI API密钥，将使用示例数据")

    # 初始化语言模型
    if zhipu_api_key:
        llm = ChatOpenAI(
            model=os.getenv("ZHIPU_MODEL", "glm-4-plus"),
            api_key=os.getenv("ZHIPU_API_KEY"),
            base_url=os.getenv("ZHIPU_BASE_URL", "https://open.bigmodel.cn/api/paas/v4")
        )
    else:
        # 如果没有API密钥，使用示例模型（仅用于演示）
        print("💡 提示：请配置ZHIPU_API_KEY环境变量以使用真实模型")
        return

    # 登录系统获取会话
    print("正在登录系统...")
    login_system = BJTSULogin(zhipu_api_key=zhipu_api_key)

    # 尝试加载已保存的会话
    if not login_system.load_session():
        print("⚠️  未找到已保存的会话，请手动登录")
        username = input("请输入学号: ").strip()
        password = input("请输入密码: ").strip()

        # 登录MIS系统
        if not login_system.login_mis_system(username, password):
            print("❌ MIS系统登录失败")
            return

        # 访问AA系统
        if not login_system.access_aa_system():
            print("❌ AA系统访问失败")
            return

        # 保存会话
        login_system.save_session()
        print("✅ 系统登录成功")

    # 获取会话
    session = login_system.get_session()

    # 创建工具
    classroom_tool = ClassroomQueryTool(session=session)
    tools = [classroom_tool]
    
    # 创建 Agent（LangChain v1.0 API）
    agent = create_agent(
        model=llm,
        tools=tools,
        system_prompt="""你是北京交通大学教室查询助手。

规则：
1. 教学楼名称：用户说什么就传什么（如"九教"、"思源西楼"）
2. 时间参数可选：不提供则自动查询当前时间
3. 时间段理解：
   - "下午" = 第5-7节
   - "上午" = 第1-2节
   - "明天" = 计算明天的日期
4. 未指定教学楼时，询问用户

直接调用工具，不要询问时间信息。"""
    )

    # 示例查询
    print("\n📝 示例查询:")
    print("1. 查一下现在思源楼哪些教室有空")
    print("2. 查一下明天思源楼的空闲教室")
    print("3. 退出")

    while True:
        user_input = input("\n请输入您的查询 (或输入 '退出' 结束): ").strip()
        if user_input.lower() in ['退出', 'quit', 'exit']:
            print("👋 再见!")
            break

        if user_input:
            try:
                response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
                
                # 提取输出
                if isinstance(response, dict) and 'messages' in response:
                    # 获取最后一条消息
                    messages = response['messages']
                    for msg in reversed(messages):
                        if hasattr(msg, 'content') and msg.content.strip():
                            print(f"\n🤖 查询结果:\n{msg.content}")
                            break
                else:
                    print(f"\n🤖 查询结果:\n{response}")
            except Exception as e:
                print(f"❌ 查询出错: {e}")
                import traceback
                traceback.print_exc()


if __name__ == "__main__":
    main()
